import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt


# 定义神经网络模型，继承自 nn.Module 类
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定义第一个卷积层，输入通道为 1，输出通道为 32，卷积核大小为 3
        self.conv1 = nn.Conv2d(1, 32, 3)
        # 定义第一个最大池化层，池化窗口大小为 2
        self.pool1 = nn.MaxPool2d(2)
        # 定义第二个卷积层，输入通道为 32，输出通道为 64，卷积核大小为 3
        self.conv2 = nn.Conv2d(32, 64, 3)
        # 定义第二个最大池化层，池化窗口大小为 2
        self.pool2 = nn.MaxPool2d(2)
        # 定义第一个全连接层，输入维度根据卷积和池化后的结果计算，输出维度为 128
        self.fc1 = nn.Linear(64 * 5 * 5, 128)
        # 定义第二个全连接层，输出维度为 10（对应 10 个类别）
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 第一个卷积层后使用 ReLU 激活函数
        x = torch.relu(self.conv1(x))
        # 第一个池化层
        x = self.pool1(x)
        # 第二个卷积层后使用 ReLU 激活函数
        x = torch.relu(self.conv2(x))
        # 第二个池化层
        x = self.pool2(x)
        # 展平操作
        x = x.view(-1, 64 * 5 * 5)
        # 第一个全连接层后使用 ReLU 激活函数
        x = torch.relu(self.fc1(x))
        # 第二个全连接层
        x = self.fc2(x)
        return x


# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.1307,), (0.3081,))  # 对数据进行归一化
])

# 加载 MNIST 训练数据集
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
# 加载 MNIST 测试数据集
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

# 创建训练数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 创建测试数据加载器
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 实例化模型
model = Net()

# 定义损失函数为交叉熵损失
loss_func = nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降，学习率为 0.01
optimizer = optim.SGD(model.parameters(), lr=0.01)


# 训练模型的函数
def train_model():
    model.train()  # 设置模型为训练模式
    for epoch in range(5):  # 进行 5 个训练轮次
        for batch_idx, (data, target) in enumerate(train_loader):  # 遍历训练数据加载器
            outputs = model(data)  # 模型的前向传播
            loss = loss_func(outputs, target)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 根据梯度更新模型参数

            if batch_idx % 100 == 0:  # 每 100 个批次打印一次损失
                print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item()}')


# 测试模型的函数
def test_model():
    # 设置模型为评估模式
    model.eval()
    # 正确预测的数量初始化
    correct = 0
    # 总样本数量初始化
    total = 0
    # 在评估时不计算梯度
    with torch.no_grad():
        # 遍历测试数据加载器
        for data, target in test_loader:
            # 模型的前向传播
            outputs = model(data)
            # 找出预测结果中概率最大的类别
            _, predicted = torch.max(outputs, 1)
            # 计算正确预测的数量，并累加到 correct 中
            correct += (predicted == target).sum().item()
            # 累加总样本数量
            total += target.size(0)
    # 计算准确率
    accuracy = correct / total
    # 打印测试集上的准确率
    print(f'Accuracy on test set: {accuracy}')


if __name__ == '__main__':
    train_model()  # 调用训练函数
    test_model()  # 调用测试函数
index = 0
for num, (image, label) in enumerate(train_loader):
    """
    遍历训练数据加载器，每次获取一个批量的图像和标签。
    num 表示当前的迭代次数。
    image 是一批图像数据。
    label 是对应的标签。
    """
    image_batch = torchvision.utils.make_grid(image, padding=2)
    """
    将一批图像组合成一个网格图像。
    设置填充为 2 像素。
    """
    plt.imshow(np.transpose(image_batch.numpy(), (1, 2, 0)), vmin=0, vmax=255)
    """
    显示图像。
    使用 numpy 的转置操作调整图像维度顺序，设置最小和最大像素值范围。
    """
    index += 1
    print(label)
    # break
    if index == 2:
        break
    else:
        plt.show()
